import cv2
import numpy as np
import os


def show_img(img, title="title"):
    cv2.imshow(title, img)
    # 关闭窗口
    cv2.waitKey()
    cv2.destroyAllWindows()


dirpath = 'temp'
imgs = [os.path.join(dirpath, img) for img in os.listdir(dirpath)]
img_a = imgs[0]


# 读取名称为 p9.png的图片
org = cv2.imread(img_a,1)
img = cv2.imread(img_a,1)
gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)


# 提取边缘
edges = cv2.Canny(gray, 30, 250, apertureSize=3)
# show_img(edges)


img_black = np.zeros(img.shape, np.uint8)
# 提取直线
COLOR = (255, 0, 0)
USE_HOUGH_P = 0
if USE_HOUGH_P:
    lines = cv2.HoughLinesP(edges, 3, 1.0 * np.pi / 180, threshold=120, minLineLength=100, maxLineGap=0)
    # lines = cv2.HoughLinesP(edges, 3, 1.0 * np.pi / 180, threshold=120)
    assert lines is not None, 'lines为空!'
    print('lines.shape:', lines.shape)

    for (x1, y1, x2, y2) in lines[:, 0]:
        print(x1, y1, ";", x2, y2)
        cv2.line(img_black, (x1, y1), (x2, y2), COLOR, 3)  # 画直线
else:
    lines = cv2.HoughLines(edges, 3, np.pi/180, 200)
    assert lines is not None, 'lines为空!'
    print('lines.shape:', lines.shape)
    for line in lines:
        for rho, theta in line:
            a = np.cos(theta)
            b = np.sin(theta)
            x0 = a * rho
            y0 = b * rho
            x1 = int(x0 + 1000 * (-b))
            y1 = int(y0 + 1000 * (a))
            x2 = int(x0 - 1000 * (-b))
            y2 = int(y0 - 1000 * (a))
            # 把直线显示在图片上
            cv2.line(img_black, (x1, y1), (x2, y2), COLOR, 2)


show_img(img_black)


# 显示原图和处理后的图像
# cv2.imshow("org",org)
# cv2.imshow("processed",img)
#
# cv2.waitKey(0)
# show_img(edges)

